# -*- coding: utf-8 -*-
"""
Created on Tue Oct 26 20:13:08 2021

@author: singularity

江榕煜的定制版YOLO网络，输出两个特征图

来自网络输出预测量形式：
(Tensor[B,3*(1+4+5),16,16],Tensor[B,3*(1+4+5),32,32])
  ↑大目标特征图              ↑小目标特征图

"""

''' 参数预设 '''

AnchorShape = [(50,44), (64,82), (82,48),
               (64,114), (108,66), (104,102)]

import torch
import torchvision
import torchvision.transforms as transforms

import torch.nn as nn
import torch.nn.functional as F

import torch.optim as optim

import matplotlib.pyplot as plt
import numpy as np

import torch.utils.tensorboard as tensorboard

import collections


'''
@定义网络
'''

def DBL(inC:int,outC:int, strid:int = 1, kerSz:int = 3):
    '''子单元：conv2D + BatchNormalization + Leaky ReLU'''
    return nn.Sequential(collections.OrderedDict([
            ("Conv",nn.Conv2d(inC, outC,    # 改变特征数量
                              kerSz,padding= int(kerSz/2),  # 不改变图大小
                              stride = strid)), # 可以通过步长做下采样
            ("BN",nn.BatchNorm2d(outC)),
            ("LReLU",nn.LeakyReLU())
            ]))

class ResUnit(nn.Module):
    '''残差子单元：1x1内核DBL+3x3内核DBL'''
    def __init__(self,inChannel:int):
        super(ResUnit, self).__init__()
        
        # NOTE：官方把特征量减少后再增多！此处先不变，有时间再来试试效果
        self.DBL_1 = DBL(inChannel,inChannel,kerSz=1)   # 1x1
        self.DBL_2 = DBL(inChannel,inChannel,kerSz=3)   # 3x3
        
    def forward(self,inFeature):
        return inFeature + self.DBL_2(self.DBL_1(inFeature))

class ResBody(nn.Module):
    '''残差块'''
    def __init__(self,inChannel:int,outChannel:int, ResN:int):
        super(ResBody,self).__init__()
        
        #self.ZP = nn.ZeroPad2d(padding=(1,1,1,1))   # 官方做的0扩充
        self.DBL = DBL(inChannel, outChannel,strid=2)   # 3x3, 步长s=2
        
        self.ResU = nn.Sequential(collections.OrderedDict(
                (str(i),ResUnit(outChannel)) for i in range(ResN)
            ) ) # 建立残差块序列
        
    def forward(self,inFeature):
        
        x = self.DBL(inFeature)
        
        return self.ResU(x)


class YOLOJ(nn.Module):
    
    def __init__(self):
        super(YOLOJ, self).__init__()
        
        '''主干特征提取线'''
        self.backbone = nn.Sequential(collections.OrderedDict([
                            ('DBL',DBL(3,32)),
                            ('Res1',ResBody(32,64,1)),
                            ('Res8',ResBody(64,256,8)),
                        ]))
        
        self.backboneTail = nn.Sequential(collections.OrderedDict([
                            ('Res4',ResBody(256,512,4) ),
                            ('DBL5_1', DBL(512,512) ),
                            ('DBL5_2', DBL(512,512) ),
                            ('DBL5_3', DBL(512,256) ),
                            ('DBL5_4', DBL(256,256) ),
# =============================================================================
#                             ('DBL5_5', DBL(256,256) ),
# =============================================================================
                        ]))
        
        '''上采样分支'''
        self.upSamp1_pre = nn.Sequential(collections.OrderedDict([
                            ('DBL',DBL(256,128)),
                            ('upsamp',nn.UpsamplingNearest2d(scale_factor=2))
                        ]))
        self.upSamp1_tail = nn.Sequential(collections.OrderedDict(
                            (str(i),DBL(384,384)) for i in range(5)
                        ))
        
        
        '''输出分支'''
        self.branch1 = nn.Sequential(
                        DBL(256,256),
                        nn.Conv2d(256, 30, 1)
                        )
        self.branch2 = nn.Sequential(
                        DBL(384,256),
                        nn.Conv2d(256, 30, 1)
                        )
        
    def forward(self,inputImg):
        '''主干特征提取'''
        featureMap1 = self.backbone(inputImg)
        featureMap2 = self.backboneTail(featureMap1)
        
        '''特征上采样'''
        upsampMap = torch.cat([featureMap1,self.upSamp1_pre(featureMap2)],dim=1)
        upsampMap =self.upSamp1_tail(upsampMap)
        
        '''输出图1'''
        out1 = self.branch1(featureMap2)
        
        '''输出图2'''
        out2 = self.branch2(upsampMap)
        
        '''激活'''
        for an in range(3):
            anPos = an*10
            # 激活 置信度和XY
            out1[:,anPos:anPos+3] = torch.sigmoid(out1[:,anPos:anPos+3])
            out2[:,anPos:anPos+3] = torch.sigmoid(out2[:,anPos:anPos+3])
            # 激活类别
            out1[:,anPos+5:anPos+10] = torch.sigmoid(out1[:,anPos+5:anPos+10])
            out2[:,anPos+5:anPos+10] = torch.sigmoid(out2[:,anPos+5:anPos+10])

        return out1,out2
    
    
    
    